import pandas as pd
import numpy as np
from sklearn.cluster import KMeans

def bin_indicators(df, bins=7, method='quantile'):
    meta_cols = ['序号', '国名En', '国名Ch', '年份']
    indicator_cols = [col for col in df.columns if col not in meta_cols]

    df_binned = df[meta_cols].copy()
    bin_meta = []

    for col in indicator_cols:
        series = df[col].dropna()

        # 初始化分箱边界
        bins_edges = None

        if method == 'quantile':
            # 等频分箱（按分位数）
            try:
                bins_edges = np.unique(np.nanpercentile(series, np.linspace(0, 100, bins + 1)))
            except Exception:
                bins_edges = np.linspace(series.min(), series.max(), bins + 1)

        elif method == 'uniform':
            # 等距分箱（每个箱间距相同）
            bins_edges = np.linspace(series.min(), series.max(), bins + 1)

        elif method == 'kmeans':
            # 聚类分箱
            reshaped = series.values.reshape(-1, 1)
            km = KMeans(n_clusters=bins, n_init=10, random_state=42)
            km.fit(reshaped)
            centers = np.sort(km.cluster_centers_.flatten())
            bins_edges = np.concatenate(([series.min() - 1e-6], (centers[:-1] + centers[1:]) / 2, [series.max() + 1e-6]))

        elif method == 'std_dev':
            # 按标准差划分（适用于近似正态分布）
            mean = series.mean()
            std = series.std()
            bins_edges = [mean + std * (i - bins // 2) for i in range(bins + 1)]
            bins_edges[0] = min(series.min(), bins_edges[0]) - 1e-6
            bins_edges[-1] = max(series.max(), bins_edges[-1]) + 1e-6
            bins_edges = sorted(bins_edges)

        else:
            raise ValueError("不支持的分箱方法，请选择 'quantile', 'uniform', 'kmeans', 'std_dev'")

        bin_labels = list(range(1, len(bins_edges)))
        df_binned[col] = pd.cut(df[col], bins=bins_edges, labels=bin_labels, include_lowest=True)

        bin_meta.append({
            'indicator': col,
            'method': method,
            'bins': bins,
            'edges': bins_edges
        })

    return df_binned, pd.DataFrame(bin_meta)
